"""
Interface definitions for adapters in continual learning.

Adapters are parameter-efficient modules that modify backbone
behavior for continual learning. This module defines interfaces
ensuring consistent behavior across different adapter implementations.
"""

from abc import abstractmethod
from typing import Any, Dict, List, Optional

import torch
from torch import nn

from continuallearning.interfaces.core.component import ComponentInterface
from continuallearning.interfaces.core.task_related_interface import (
    TaskIrrelevantInterface,
)
from continuallearning.interfaces.models.pefts.hooks.hook_manager_interface import (
    HookManagerInterface,
)
from continuallearning.interfaces.types import AdapterOutput


class AdapterInterface(ComponentInterface, nn.Module, TaskIrrelevantInterface):
    """
    Interface for model adapter implementations.

    Adapters are small, trainable modules that enable parameter-efficient
    fine-tuning of backbone models. They can modify either the backbone itself
    or its outputs to facilitate continual learning.
    """

    @abstractmethod
    def forward(
        self, x: torch.Tensor, task_ids: Optional[List[int]] = None, **kwargs
    ) -> AdapterOutput:
        """
        Process features through the adapter.

        Args:
            x: Input tensor or features depending on adapter_type
            task_ids: Optional task identifiers, defaults to current_task if None
            **kwargs: Additional arguments

        Returns:
            AdapterOutput: Standardized adapter output format
        """
        pass

    @abstractmethod
    def prepare_task(self, task_id: int, **kwargs) -> None:
        pass

    @abstractmethod
    def register_task(
        self, task_id: int, config: Optional[Dict[str, Any]] = None
    ) -> None:
        """
        Register a new task with this hook.

        Args:
            task_id: Identifier for the new task
            config: Optional task-specific configuration
        """
        pass

    @abstractmethod
    def set_task(self, task_id: int) -> None:
        """
        Configure the model for a specific task.

        Args:
            task_id: Task identifier
        """
        self._current_task = task_id

    @abstractmethod
    def set_trainable_tasks(self, task_ids: int | List[int], **kwargs) -> None:
        """
        Prepare this hook for a specific task.

        Args:
            task_id: Identifier for the task to prepare
            **kwargs: Additional preparation parameters
        """
        pass

    @property
    @abstractmethod
    def current_task(self) -> int:
        """Get the current task ID."""
        pass

    @property
    @abstractmethod
    def registered_tasks(self) -> List[int]:
        """Get list of registered task IDs."""
        pass


class HookBasedAdapterInterface(AdapterInterface):
    """
    Interface for hook-based adapters.

    Hook-based adapters use PyTorch hooks instead of directly modifying
    model architecture, providing better separation between components.
    They can also function as context managers for automatic hook management.
    """

    @abstractmethod
    def attach(
        self,
    ) -> HookManagerInterface:
        """
        Attach this adapter to backbone using hooks.
        Returns:
            HookManager: Hook manager instance for the attached hooks
        """
        pass

    @abstractmethod
    def detach(self) -> None:
        """Detach this adapter from the model, removing all hooks."""
        pass

    @abstractmethod
    def __enter__(self) -> "HookBasedAdapterInterface":
        """
        Enter context manager, attaching hooks.

        Returns:
            HookBasedAdapterInterface: Self reference for context manager usage
        """
        pass

    @abstractmethod
    def __exit__(
        self, exc_type: Optional[Any], exc_val: Optional[Any], exc_tb: Optional[Any]
    ) -> None:
        """
        Exit context manager, removing hooks.

        Args:
            exc_type: Exception type if an exception was raised
            exc_val: Exception value if an exception was raised
            exc_tb: Exception traceback if an exception was raised
        """
        pass
